# encoding=utf-8

# 这个文件由pnnx程序生成
from v1_onnx_pnnx import Model
import torch
import torch.nn.functional as F
import time
import numpy
import cv2


model = Model()

model.eval()
torch.manual_seed(0)

# load data
sample_image = "../resource/raw_resize/1.jpg"
start_time = time.perf_counter()
mat = cv2.imread(sample_image)
mat = mat.transpose((2, 0, 1))
mat = mat.astype(numpy.float32)
for c in range(mat.shape[0]):
    mat[c] = (mat[c] - mat[c].mean()) / (mat[c].std() + 1e-8)
mat = numpy.expand_dims(mat, 0)
# infer
with torch.no_grad():
    x = torch.from_numpy(mat)
    pred = F.softmax(model(x), 1)
    pred += torch.flip(F.softmax(model(torch.flip(x, (3,))), 1), (3,))
    pred += torch.flip(F.softmax(model(torch.flip(x, (2,))), 1), (2,))
    pred += torch.flip(F.softmax(model(torch.flip(x, (3, 2))), 1), (3, 2))
    pred /= 4
    pred = pred.argmax(1)
    # save
    pred = pred.cpu().numpy()[0]
    cv2.imwrite("out_pred_pnnx.png", pred * 255)
end_time = time.perf_counter()
print("time cost: ", end_time - start_time, "s")